import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

# ### 操作目标不同
# - all_gather ： 目标是 收集所有进程的原始数据 。每个进程将自己的本地数据发送给所有其他进程，最终每个进程都持有所有进程的原始数据（即全局完整数据）。 例如，在 all_gather.py 中，4 个进程分别生成 [1,2] 、 [3,4] 、 [5,6] 、 [7,8] ， all_gather 后每个进程的 gather_list 会包含这 4 组原始数据。
# - all_reduce ： 目标是 对所有进程的数据进行归约（如求和、取最大值等） ，并将归约后的结果同步给所有进程。最终每个进程得到的是归约后的值，而非原始数据。 例如，在之前的 all_reduce_train.py 中，每个进程计算梯度后， all_reduce 会将所有梯度相加，最终每个进程得到梯度的总和（后续再取平均）。
# ### 结果形式不同
# - all_gather 的结果 ： 每个进程保留所有原始数据（如 gather_list 包含所有进程的原始张量）。 适用场景：需要全局原始数据的场景（如分布式数据拼接、多机联合特征计算等）。
# - all_reduce 的结果 ： 每个进程仅保留归约后的值（如所有梯度的和、所有损失的最大值等）。 适用场景：需要全局统计结果的场景（如梯度平均、全局损失汇总等）。
# ### 效率差异（为何 all_gather 不能替代 all_reduce）
# 虽然理论上可以通过 all_gather 收集所有数据后手动计算归约（如求和），但实际中 all_reduce 更高效：

# - 数据传输量 ： all_gather 需要每个进程发送完整的原始数据（总传输量为 N * 数据量 ， N 为进程数）；而 all_reduce 在传输过程中直接归约（如累加），传输量仅为 数据量 （每个进程只需发送一次数据，后续通过归约逐步合并）。
# - 计算开销 ： all_reduce 的归约操作由底层通信库（如 gloo / nccl ）优化实现（可能利用硬件加速），比手动收集后计算更高效。

def run(rank_id, size):
    tensor = torch.arange(2, dtype=torch.int64) + 1 + 2 * rank_id
    print('before reudce',' Rank ', rank_id, ' has data ', tensor)
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
    print('after reudce',' Rank ', rank_id, ' has data ', tensor)


def init_process(rank_id, size, fn, backend='gloo'):
    """ Initialize the distributed environment. """
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '29500'
    dist.init_process_group(backend, rank=rank_id, world_size=size)
    fn(rank_id, size)


if __name__ == "__main__":
    size = 4
    processes = []
    mp.set_start_method("spawn")
    for rank in range(size):
        p = mp.Process(target=init_process, args=(rank, size, run))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()
